# src/matrix/matirx_utility.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*

import scipy.sparse as sparse

def read_data(matrix_path):
    rptr = []
    col = []
    value = []
    i     = 1
    n     = 0
    nnz   = 0
    row_first = 0
    col_first = 0
    print("read matrix: \n" +matrix_path)
    for line in open(matrix_path, "r" ):
        line = line.strip("\n")
        if(i == 1):
            n = int(line)
        elif(i > 1 and i <= n + 2):
            if(i == 2):
                row_first = int(line)
            elif(i == n + 2):
                nnz = int(line) - row_first
            rptr.append(int(line) - row_first)
        elif(i > n + 2 and i < n + 3 + nnz):
            if(i == n + 3):
                col_first =int(line)
            col.append(int(line) - col_first)
        else:
            value.append(float(line))
            # print("%.15e" %(float(line)))
        i = i + 1
    print("read matrix success: n is {}, nnz is {}".format(n,nnz))

    csr_mat = sparse.csr_matrix((value, col, rptr,), shape=(n, n))
    return csr_mat

def save_matrix(mat, save_path):
    csr_mat = mat.tocsr()

    n    = csr_mat.shape[0]
    nnz  = csr_mat.nnz
    rptr = csr_mat.indptr
    col  = csr_mat.indices
    value  = csr_mat.data

    with open(save_path, "w+" ) as fo:
        fo.write(str(n)+"\n")
        for i in range(0, n + 1):
            # fo.write(str(rptr[i])+"\n")
            fo.write("%d" %(rptr[i])+"\n")
        for i in range(0, nnz):
            # fo.write(str(col[i])+"\n")
            fo.write("%d" %(col[i])+"\n")
        for i in range(0, nnz):
            # fo.write(str(value[i])+"\n")
            fo.write("%.15e" %(value[i])+"\n")

def test():
    return

if __name__=="__main__":
    test()
